#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import os
import time
import re
import datetime

HOST = 'localhost'
PORT = 27017
DB_NAME = 'pubmed'
COLLECTION_NAME = 'metalist'


def GetConfig():
    if os.path.exists("../config.ini"):
        import ConfigParser
        global HOST
        global PORT
        global DB_NAME
        global COLLECTION_NAME
        config = ConfigParser.ConfigParser()
        config.read("../config.ini")
        HOST = config.get("mongodb", "host")
        PORT = int(config.get("mongodb", "port"))
        DB_NAME = config.get("mongodb", "db")
        COLLECTION_NAME = config.get("mongodb", "collection")


def GetPMID(line):
    pattern = re.compile("<.*>(.*)<.*>", re.I)
    m = re.search(pattern, line)
    if m:
        return m.group(1)
    else:
        return m


# Use regex to analyze PMID lines
def Extract_PubMed_IDs_From_XML(filename):
    infile = open(filename, "r")
    outfile = open(filename[:-3] + "ids", "w")
    result = []
    for line in infile:
        line = line.rstrip()
        if line.startswith("        <PMID"):
            pmid_string = GetPMID(line)
            if pmid_string.isdigit() and pmid_string != None:
                result.append(pmid_string)
                outfile.write(pmid_string + "\n")
                outfile.flush()
    outfile.close()
    infile.close()
    return result


# Total Analyze
# def Extract_PubMed_IDs_From_XML(filename):
#     from bs4 import BeautifulSoup
#     infile = open(filename, "r")
#     outfile = open(filename[:-3] + "ids", "w")
#     result = []
#     current_string = ""
#     for line in infile:
#         line = line.rstrip()
#         current_string += line
#         if (line.find("</PubmedArticle>") != -1) or (line.find("</PubmedBookArticle>") != -1):
#             soup = BeautifulSoup(current_string, "xml")
#             pmids = soup.find_all("PMID")
#             for pmid in pmids:
#                 parent_name = pmid.parent.name
#                 if (parent_name == "MedlineCitation") or (parent_name == "BookDocument"):
#                     pmid_string = pmid.string
#                     if pmid_string.isdigit():
#                         result.append(pmid_string)
#                         outfile.write(pmid_string + "\n")
#                         outfile.flush()
#             current_string = ""
#     outfile.close()
#     infile.close()
#     return result

if __name__ == '__main__':
    print("Loading config file ... "),
    time.sleep(1)
    GetConfig()
    print("complete.")
    time.sleep(1)

    print("Initiating mongodb connection ... "),
    time.sleep(1)
    connection = MongoClient(HOST, PORT)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]
    print("complete.")
    time.sleep(1)

    print("Loading PubMed update list ... "),
    time.sleep(1)
    updatelist = {}
    dateString = os.path.abspath(".").split("_")[-1]
    updateListFilename = "pmid_todo.update." + dateString + ".txt"
    if os.path.getsize(updateListFilename) != 0:
        updatefile = open(updateListFilename, "r")
        for line in updatefile:
            line = line.rstrip()
            splitline = line.split("\t")
            updatelist[int(splitline[0])] = splitline[1]
        updatefile.close()
    print("complete.")
    time.sleep(1)

    print("Loading PubMed xml files ... "),
    time.sleep(1)
    filelist = os.listdir(".")
    print("complete.")
    time.sleep(1)

    print("Start anlyzing xml files ... ")
    time.sleep(1)
    file_count = 0
    pmid_count = 0
    is_update_file = False
    for filename in filelist:
        if filename.endswith(".xml"):
            filepath = os.path.dirname(os.path.abspath(filename)).split("/")[-1] + "/" + filename
            file_count += 1
            if filename.find("update") != -1:
                is_update_file = True
            else:
                is_update_file = False
            print("Start analyzing " + filename + " ... "),
            ids = Extract_PubMed_IDs_From_XML(filename)
            pmid_count += len(ids)
            for id in ids:
                pmid = int(id)
                collection.update_one({"pmid": pmid}, {"$set": {"xml": filepath}})
                if is_update_file:
                    if updatelist.has_key(pmid):
                        updatelist.pop(pmid)
            print("complete. Find " + str(len(ids)) + " PMIDs.")
            time.sleep(1)

    print("Reverse no xml update records ... "),
    time.sleep(1)
    if len(updatelist) != 0:
        for key in updatelist:
            if updatelist[key] == "None":
                collection.update_one({"pmid": pmid}, {"$unset": {"pm_date": True}})
            else:
                origin_datetime = datetime.datetime.strptime(updatelist[key], "%Y%m%d%H")
                collection.update_one({"pmid": pmid}, {"$set": {"pm_date": origin_datetime}})
    print(str(len(updatelist)) + " records reversed.")
    time.sleep(1)

    connection.close()
    print("All xml files analyzed. " + str(file_count) + " files and " + str(pmid_count) + " PMIDs found.")
